# web_1px_pen.py  ——  笔刷 1 px 大屏版
from __future__ import annotations
import os
import pickle
from typing import Any, Tuple
import numpy as np

try:
    import gradio as gr
except Exception:
    gr = None

try:
    from PIL import Image
except Exception:
    Image = None


def load_knn_model() -> Any:
    model_path = "best_knn_model.pkl"   # 确保 k=1
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"请先把 k=1 模型保存为 {model_path}")
    with open(model_path, "rb") as f:
        model = pickle.load(f)
    print(f"[OK] 加载 k=1 模型：{type(model).__name__}")
    return model


MODEL = load_knn_model()


def preprocess_image_for_digits(img: Any) -> Tuple[np.ndarray, np.ndarray]:
    if img is None:
        raise ValueError("输入图像为空")
    if isinstance(img, np.ndarray):
        img = Image.fromarray(img.astype(np.uint8))
    img = img.convert("L").resize((8, 8), Image.LANCZOS)
    arr = np.asarray(img, dtype=np.float32)
    arr = 255.0 - arr
    arr = (arr / 255.0) * 16.0
    vec = arr.flatten()
    processed_img = (vec.reshape(8, 8) / 16.0 * 255).astype(np.uint8)
    Image.fromarray(processed_img).save("debug_processed.png")
    print("64 维向量：", np.round(vec, 2))
    return vec, processed_img


def predict_from_canvas(img: Any) -> Tuple[str, np.ndarray]:
    try:
        x, processed_img = preprocess_image_for_digits(img)
        pred = int(MODEL.predict([x])[0])
        
        # Get prediction probabilities if available
        probabilities = ""
        if hasattr(MODEL, 'predict_proba'):
            probs = MODEL.predict_proba([x])[0]
            # Get top 3 predictions
            top3_idx = np.argsort(probs)[::-1][:3]
            probabilities = "\n\nTop 预测:\n" + "\n".join([
                f"{i+1}. {idx}: {probs[idx]:.2%}" 
                for i, idx in enumerate(top3_idx)
            ])
        
        result = f"预测结果：{pred}{probabilities}"
        return result, processed_img
    except Exception as e:
        return f"预测错误：{str(e)}", np.zeros((8, 8), dtype=np.uint8)


def clear_canvas() -> Tuple[None, str, np.ndarray]:
    """Clear the canvas and reset outputs"""
    return None, "请在画布上绘制一个数字", np.zeros((8, 8), dtype=np.uint8)


def create_and_launch_app():
    if gr is None:
        raise RuntimeError("请 pip install gradio==3.50.2")

    with gr.Blocks(title="KNN 手写数字识别") as iface:
        gr.Markdown("# KNN 手写数字识别（1 px 笔刷版）")
        gr.Markdown("在画布上绘制一个数字，然后点击 Submit 查看预测结果")
        
        with gr.Row():
            with gr.Column():
                canvas = gr.Image(
                    source="canvas",
                    type="numpy",
                    image_mode="L",
                    shape=(40, 40),
                    label="在白画布内写一个小而细的数字（0-9）"
                )
                with gr.Row():
                    submit_btn = gr.Button("Submit", variant="primary")
                    clear_btn = gr.Button("Clear")
            
            with gr.Column():
                output_text = gr.Textbox(label="预测结果", value="请在画布上绘制一个数字")
                processed_image = gr.Image(label="预处理后的 8×8 图像", shape=(160, 160))
        
        gr.Markdown("### 提示")
        gr.Markdown("- debug_processed.png 文件会保存预处理后的图像\n- 绘制时请尽量在画布中央书写数字")
        
        # Event handlers
        submit_btn.click(
            fn=predict_from_canvas,
            inputs=canvas,
            outputs=[output_text, processed_image]
        )
        
        clear_btn.click(
            fn=clear_canvas,
            inputs=None,
            outputs=[canvas, output_text, processed_image]
        )
        
        # Also clear on image change (optional)
        canvas.change(
            fn=lambda: ("请在画布上绘制一个数字", np.zeros((8, 8), dtype=np.uint8)),
            inputs=None,
            outputs=[output_text, processed_image]
        )
    
    iface.launch(share=False, server_name="127.0.0.1", server_port=7860)


if __name__ == "__main__":
    create_and_launch_app()